#!/usr/bin/env python
# -*- encoding: utf-8 -*-
"""
@File    :   visualize.py
@Time    :   2021/12/01 16:13:20
@Author  :   glx 
@Version :   1.0
@Contact :   18095542g@connect.polyu.hk
@Desc    :   visualize data including 3D trajectory / Frequency / velocity / Azimuth
"""

__all__ = ["visualize_data", "visualize_single", "do_visualize", "visualize_smooth"]


# here put the import lib
import matplotlib
from matplotlib import lines
from pathlib import Path

# matplotlib.use("Agg")
matplotlib.use("TkAgg")
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import os
from copy import deepcopy
from typing import List


def visualize_data(data, save_data_list, save_folder, FIGSIZE):
    """
    visualize data
    data: Dict()
    save_data_list: List()
    """
    os.makedirs(save_folder, exist_ok=True)
    for target in save_data_list:
        plt.figure(figsize=FIGSIZE)
        for data_type in data.keys():
            if data_type == "Traj":
                continue
            plt.plot(data[data_type][target], label=data_type, linewidth=1)

        plt.xlabel("time", fontsize=32)
        plt.ylabel("value", fontsize=32)
        plt.xticks(fontsize=32)
        plt.yticks(fontsize=32)
        plt.legend(prop={"family": "Times New Roman", "size": 16})
        plt.savefig(os.path.join(save_folder, target) + ".png", dpi=300)
        plt.close("all")


def visualize_single(
    data_list, data_name, save_data_list, folder_name, three_D, FIGSIZE, SAVE_FIG=True
):
    """draw single data
    data_list: List(Dict())
    data_ame: List(str)
    save_data_list: List(str)
    folder_name: str
    """
    os.makedirs(folder_name, exist_ok=True)

    if len(data_list) != len(data_name):
        raise ValueError("data_list and data_name should have same length")

    for target in save_data_list:
        # print(data.keys())
        fig = plt.figure("graph", figsize=FIGSIZE)

        num = 0
        # visual_data = copy.deepcopy(data[target])

        if three_D:  # 3d data
            ax = fig.add_subplot(111, projection="3d")  # 创建一个三维的绘图工程
            try:
                for data in data_list:
                    visual_data = deepcopy(data[target])
                    ax.plot(
                        visual_data["pos_x"],
                        visual_data["pos_y"],
                        visual_data["pos_z"],
                        label=target,
                    )
                ax.set_xlabel("x", fontsize=32)
                ax.set_ylabel("y", fontsize=32)
                ax.set_zlabel("z", fontsize=32)
            except ValueError as e:
                print("not 3D data")
                raise e
        else:  # 2d data
            ax = fig.add_subplot()
            for data in data_list:
                # 画点
                if data_name[num].endswith("point"):
                    ax.scatter(
                        data[target].index,
                        data[target].values,
                        label=data_name[num],
                        color="red",
                        marker="x",
                    )
                    # TO DO
                    pass

                # 画线
                else:
                    # 平滑曲线宽度为2
                    if data_name[num] == "smooth":
                        width = 2
                    else:
                        width = 1
                    ax.plot(
                        data[target],
                        label=target + "_" + data_name[num],
                        linewidth=width,
                    )

                num += 1
            ax.set_xlabel("time", fontsize=32)
            ax.set_ylabel("value", fontsize=32)
            ax.legend()
            plt.xticks(fontsize=32)
            plt.yticks(fontsize=32)
        if SAVE_FIG:
            plt.savefig(os.path.join(folder_name, target) + ".png")
        plt.close("graph")


def do_visualize(data, smooth_data, FIGSIZE):
    """do visualize and save in folder pictures"""
    save_folder = r".\picture"
    # 分项文件夹
    traj_folder = r".\traj"
    freq_folder = r".\freq"
    total_folder = r".\total"

    # 可视化数据
    save_data_list = ["aircraft", "ship", "station", "car"]
    # 轨迹
    visualize_single(
        [data["Traj"]],
        ["origin"],
        save_data_list,
        os.path.join(save_folder, traj_folder),
        three_D=True,
        FIGSIZE=FIGSIZE,
    )

    # 所有图像
    visualize_data(
        data, save_data_list, os.path.join(save_folder, total_folder), FIGSIZE
    )

    # 可视化平滑数据
    visualize_smooth(data, smooth_data, FIGSIZE)
    # visualize_single([data["Freq"], smooth_data], ["orgin","smooth",],  save_data_list,  os.path.join(save_folder, freq_folder), three_D = False)
    print("visualize done".center(100, "="))


def visualize_smooth(data, smooth_data, FIGSIZE):
    save_folder = r".\picture"
    # 分项文件夹
    freq_folder = r".\freq"
    # 可视化数据
    save_data_list = ["aircraft", "ship", "station", "car"]
    visualize_single(
        [data["Freq"], smooth_data],
        [
            "orgin",
            "smooth",
        ],
        save_data_list,
        os.path.join(save_folder, freq_folder),
        three_D=False,
        FIGSIZE=FIGSIZE,
    )
    print("visualize smooth data done".center(100, "="))


def visualize_diff(
    data,
    smooth_data,
    turbulence_point_index,
    diff_data,
    ERROR_N,
    upper_bound,
    lower_bound,
    if_debug,
    FIGSIZE,
):
    """visualize turbulence result"""

    save_folder = r".\picture"
    # 分项文件夹
    diff_data_folder = r".\diff"
    turbulence_data_folder = r".\turbulence"

    save_data_list = ["aircraft", "ship", "station", "car"]

    turbulence_point = {}
    for key in turbulence_point_index.keys():
        turbulence_point[key] = smooth_data[key].iloc[turbulence_point_index[key]]

    # 可视化突变点
    if if_debug:
        visualize_single(
            [
                diff_data,
                smooth_data,
                turbulence_point,
                upper_bound,
                lower_bound,
            ],
            [
                "diff_data",
                "Freq",
                "diff_point",
                "Velo",
                "Azimuth",
                "upper_diff",
                "lower_diff",
            ],
            save_data_list,
            os.path.join(save_folder, turbulence_data_folder + "_debug", str(ERROR_N)),
            three_D=False,
            FIGSIZE=FIGSIZE,
        )
    else:
        visualize_single(
            [smooth_data, turbulence_point, data["Velo"], data["Azimuth"]],
            ["Freq", "diff_point", "Velo", "Azimuth"],
            save_data_list,
            os.path.join(save_folder, turbulence_data_folder, str(ERROR_N)),
            three_D=False,
            FIGSIZE=FIGSIZE,
        )

    print("visualize diff done".center(100, "="))


def visualize_ADF_result(test_result, data, step, Freq_range, if_debug, FIGSIZE):
    """visualize ADF result"""
    save_folder = r".\picture"
    # 分项文件夹
    stable_data_folder = r".\stable_test"

    save_data_list = ["aircraft", "ship", "station", "car"]

    os.makedirs(save_folder, exist_ok=True)

    # 对每个目标画图
    for key in test_result.keys():
        plt.figure(figsize=FIGSIZE)
        plt.plot(data["Azimuth"][key], label="Azimuth")
        plt.plot(data["Freq"][key], label="Freq", color="grey")
        plt.plot(data["Velo"][key], label="Velo")

        for i in test_result[key]:
            if i[0]:  # 如果是平稳态:
                # stable_.append(data["Freq"][key].iloc[i[1]:i[2]])
                plt.plot(data["Freq"][key].iloc[i[1] : i[2]], color="green")
        # plt.plot(stable_, label = "stable", color = "green")
        plt.legend()
        if if_debug:
            os.makedirs(
                os.path.join(
                    save_folder,
                    stable_data_folder + "_debug",
                    str(step) + "_" + str(Freq_range),
                ),
                exist_ok=True,
            )
            plt.savefig(
                os.path.join(
                    save_folder,
                    stable_data_folder + "_debug",
                    str(step) + "_" + str(Freq_range),
                    str(key),
                )
                + ".png"
            )
        else:
            os.makedirs(os.path.join(save_folder, stable_data_folder), exist_ok=True)
            plt.savefig(
                os.path.join(save_folder, stable_data_folder, str(key)) + ".png"
            )
        plt.close("all")
    print("visualize ADF result done".center(100, "="))


def visualize_fast_result(test_result, data, step, Freq_range, if_debug, FIGSIZE):
    """visualize ADF result"""
    save_folder = r".\picture"
    # 分项文件夹
    stable_data_folder = r".\fast_test"

    os.makedirs(save_folder, exist_ok=True)

    # 对每个目标画图
    for key in test_result.keys():
        plt.figure(figsize=FIGSIZE)
        # plt.plot(data["Azimuth"][key], label="Azimuth")
        # plt.plot(data["Freq"][key], label="Freq", color="grey")
        # plt.plot(data["Velo"][key], label="Velo")
        stable_ = []
        for i in test_result[key]:
            if i[0]:  # 如果是平稳态:
                # stable_.append(data["Freq"][key].iloc[i[1]:i[2]])
                plt.plot(pd.Series(data["signal"]).iloc[i[1] : i[2]], color="orange")
        # plt.plot(stable_, label = "stable", color = "green")
        plt.legend()
        if if_debug:
            os.makedirs(
                os.path.join(
                    save_folder,
                    stable_data_folder + "_debug",
                    str(step) + "_" + str(Freq_range),
                ),
                exist_ok=True,
            )
            plt.savefig(
                os.path.join(
                    save_folder,
                    stable_data_folder + "_debug",
                    str(step) + "_" + str(Freq_range),
                    str(key),
                )
                + ".png"
            )
        else:
            os.makedirs(os.path.join(save_folder, stable_data_folder), exist_ok=True)
            plt.savefig(
                os.path.join(save_folder, stable_data_folder, str(key)) + ".png"
            )
        plt.close("all")
    print("visualize fast result done".center(100, "="))


def visualize_total_result(
    data,
    smooth_data,
    diff_data,
    turbulence_point_index,
    upper_bound,
    lower_bound,
    stable_test_result,
    fast_test_result,
    static_test_result,
    FAST_STEP,
    SLOW_STEP,
    FAST_FREQ_RANGE,
    SLOW_FREQ_RANGE,
    TAN_RANGE,
    ST,
    DEBUG,
    FIGSIZE,
    save_dir,
):
    """可视化所有结果"""
    save_dir = Path(save_dir) if not isinstance(save_dir, Path) else save_dir
    parent_folder = save_dir.parent

    # res 如果是一个list中有多个元素，则把所有元素合并为一个list
    # res = [i for i in res for i in i]
    stable_test_result = [i for i in stable_test_result for i in i]
    fast_test_result = [i for i in fast_test_result for i in i]
    static_test_result = [i for i in static_test_result for i in i]

    # 对每个目标画图

    turbulence_point = {}

    plt.figure(figsize=FIGSIZE)
    turbulence_point = smooth_data[["Freq", "ts_range"]].iloc[
        turbulence_point_index["Freq"]
    ]

    # TODO : FIX BUG
    plt.plot(smooth_data["ts_range"], smooth_data["Freq"], label="freq")

    # 画突变点
    plt.scatter(
        turbulence_point["ts_range"].values,
        turbulence_point["Freq"].values,
        label="motor moving",
        color="red",
        marker="x",
    )

    for i in fast_test_result:
        if i[0]:  # 如果快速运动:
            # stable_.append(data["Freq"][key].iloc[i[1]:i[2]])
            plt.plot(
                smooth_data["ts_range"][i[1] : i[2] + 1],
                smooth_data["Freq"][i[1] : i[2] + 1],
                color="green",
                label="fast",
            )

    for i in stable_test_result:
        if i[0]:  # 如果是平稳态:
            plt.plot(
                smooth_data["ts_range"].iloc[i[1] : i[2] + 1],
                smooth_data["Freq"].iloc[i[1] : i[2] + 1],
                color="orange",
                label="slow",
            )
        # print(stable_test_result[-1])

    for i in static_test_result:
        if i[0]:  # 如果是静止态:
            plt.plot(
                smooth_data["ts_range"].iloc[i[1] : i[2] + 1],
                smooth_data["Freq"].iloc[i[1] : i[2] + 1],
                color="pink",
                label="static",
            )
        # print(static_test_result[-1])
        # tl = smooth_data["ts_range"].iloc[static_test_result[-1][1]]
    # 画出tan_range确定的斜率范围
    if DEBUG:
        for i in [TAN_RANGE[0], TAN_RANGE[1], ST]:

            for interval in fast_test_result:
                # if interval[1] < len(smooth_data["ts_range"]):
                # print(len(smooth_data))
                x_0 = smooth_data["ts_range"].iloc[interval[1]]
                y_0 = smooth_data["Freq"].iloc[interval[1]]
                x = smooth_data["ts_range"].iloc[interval[1] : interval[2]]
                y = np.tan(np.deg2rad(i)) * (x - x_0) + y_0
                plt.plot(x, y, color="black", linestyle="--")
    # else:
    #     print(
    #         f"Warning: index {interval[1]} is out of bounds. Skipping this interval."
    #     )
    #     continue

    handles, labels = plt.gca().get_legend_handles_labels()
    by_label = dict(zip(labels, handles))
    plt.legend(by_label.values(), by_label.keys())

    parent_folder.mkdir(exist_ok=True, parents=True)
    # plt.savefig(
    #     os.path.join(save_folder, stable_data_folder + "_debug", str(key))
    #     + ".png"
    # )
    plt.savefig(str(parent_folder) + "\\" + str(save_dir.stem) + ".jpg")
    # plt.show()
    plt.close("all")
    print("visualized all ".center(100, "="))


if __name__ == "__main__":
    from get_matlab_data import get_data
    from smooth_data import do_smooth
    from diff_test import do_diff_test
    from ADF_test import do_ADF_test
    from fast_test import do_fast_test

    data = get_data()
    smooth_data = do_smooth(data, sigma=0)

    # 一阶diff
    # diff_data, turbulence_point_index, upper_bound, lower_bound = do_diff_test(smooth_data, 1)

    # visualize_diff(data,smooth_data, turbulence_point_index, diff_data, 3, upper_bound, lower_bound, if_debug=True)
    # do_visualize(data, smooth_data)

    # 平稳检测
    # 实验
    FIGSIZE = (25, 12)
    STEP = [100]  # 步长

    TAN_RANGE = [(10, 30)]
    FREQ_RANGE = [10]  # 极差阈值
    ST = 10  # 静止阈值
    # for step in STEP:
    #     for Freq_range in FREQ_RANGE:

    #         test_result, stable_data = do_ADF_test(smooth_data, step, Freq_range, ST)
    #         visualize_ADF_result(test_result, data, step, Freq_range, if_debug=True)

    # 快速检测
    # for step in STEP:
    #     for Freq_range in FREQ_RANGE:
    #         test_result, stable_data = do_fast_test(smooth_data, step, Freq_range, ST)
    #         visualize_ADF_result(test_result, data, step, Freq_range, if_debug=True)

    test_result, stable_data = do_fast_test(
        smooth_data, step=10, upper_rad=90, lower_rad=10, Freq_range=10
    )
    visualize_fast_result(
        test_result, data, step=100, Freq_range=10, if_debug=True, FIGSIZE=FIGSIZE
    )
